# -*- coding: utf-8 -*-
# 全样本预测 + 指标计算

import torch, json, glob, config, os
from transformers import BertForSequenceClassification, BertTokenizerFast
from tqdm import tqdm

tok = BertTokenizerFast.from_pretrained("models/finbert-risk")
model = BertForSequenceClassification.from_pretrained("models/finbert-risk")
model.eval()

def predict_sent(sent):
    inputs = tok(sent, return_tensors="pt", truncation=True, padding=True, max_length=128)
    with torch.no_grad():
        logits = model(**inputs).logits
    probs = torch.softmax(logits, dim=1)[0]
    return {"exposure": probs[2].item(), "defense": probs[1].item(), "neutral": probs[0].item()}

results = []
for js in tqdm.tqdm(glob.glob("data/filtered/*.json")):
    ticker_year = os.path.basename(js).replace(".json", "")
    sentences = json.load(open(js, encoding="utf-8"))
    e_probs, d_probs = [], []
    for s in sentences:
        p = predict_sent(s)
        if p["exposure"] > 0.5:
            e_probs.append(p["exposure"])
        if p["defense"] > 0.5:
            d_probs.append(p["defense"])
    exposure = max(e_probs) if e_probs else 0
    defense  = sum(d_probs)/len(d_probs) if d_probs else 0
    risk_score = max(0, exposure - defense)
    results.append({"firm": ticker_year, "risk": risk_score})

import pandas as pd
pd.DataFrame(results).to_csv("outputs/risk_metrics.csv", index=False)
print("✅ 风险指标已写入 outputs/risk_metrics.csv")

